from re import findall
import os, sys

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from monitoring import get_mysql_data_management, get_mysql_hive
from my_utils.timestamp_handler import second2str
from my_utils.hive import Hive

# 数据存储路径
MYSQL_TB1 = 'hive_table_status'
MYSQL_TB2 = 'hive_partition_status'
# 查询HIVE表
SQL = '''
SELECT
    b.`NAME`     AS db_name,  -- 库名
    t.`TBL_NAME` AS tb_name   -- 表名
FROM `DBS` b
INNER JOIN `TBLS` t ON b.`DB_ID`=t.`DB_ID`
WHERE 1=1
'''
HIVE_DB = ''  # 根据库来筛选HIVE表，默认不筛选
EXTERNAL_TABLE = False   # 是否筛选外部表


class Monitor:
    def __init__(self):
        self.mysql_hive = get_mysql_hive()
        self.mysql_db = get_mysql_data_management()
        self.hive = Hive()
        self.hive.set_user('hdfs', False)
        self.hive_result = ''

    def find_all(self, pattern):
        """正则表达式查找HIVE表各参数"""
        a = findall(pattern, self.hive_result)
        if a:
            a = a[0]
            if a.isdigit():
                return int(a)
            return a
        return None

    def select_hive_tables(self):
        sql = SQL
        if HIVE_DB:
            sql += ' AND b.`NAME`="{hive_db}"'.format(hive_db=HIVE_DB)
        if EXTERNAL_TABLE is True:
            sql += ' AND `TBL_TYPE`="EXTERNAL_TABLE"'
        for dt in self.mysql_hive.fetchone_dt(sql):
            yield '{}.{}'.format(dt['db_name'], dt['tb_name'])

    def desc_formatted(self, table):
        sql = 'DESC FORMATTED {};'.format(table)
        self.hive_result = self.hive.e(sql)
        return {
            'table_name': table,
            'num_partitions': self.find_all(r'numPartitions\s+([0-9]+)'),
            'raw_data_size': self.find_all(r'rawDataSize\s+([0-9]+)'),
            'total_size': self.find_all(r'totalSize\s+([0-9]+)'),
            'num_rows': self.find_all(r'numRows\s+([0-9]+)'),
            'num_files': self.find_all(r'numFiles\s+([0-9]+)'),
        }

    def desc_formatted_partition(self, table, partition=None):
        partition = partition or self.hive.ymd
        sql = 'DESC FORMATTED {} PARTITION(ymd="{}");'.format(table, partition)
        self.hive_result = self.hive.e(sql)
        return {
            'table_name': table,
            'partition_value': partition,
            'raw_data_size': self.find_all(r'rawDataSize\s+([0-9]+)'),
            'total_size': self.find_all(r'totalSize\s+([0-9]+)'),
            'num_rows': self.find_all(r'numRows\s+([0-9]+)'),
            'num_files': self.find_all(r'numFiles\s+([0-9]+)'),
            'transient_last_ddl_time': second2str(self.find_all(r'transient_lastDdlTime\s+([0-9]+)')),
        }

    def main(self):
        for tb in self.select_hive_tables():
            try:
                self.mysql_db.replace(self.desc_formatted_partition(tb), MYSQL_TB2)
            except Exception as e:
                self.mysql_db.write_log(e, '/tmp/log-hive_table_status')
            try:
                self.mysql_db.replace(self.desc_formatted(tb), MYSQL_TB1)
            except Exception as e:
                self.mysql_db.write_log(e, '/tmp/log-hive_partition_status')


if __name__ == '__main__':
    Monitor().main()
